package com.codejiwei.core.graphx

import org.apache.spark.SparkContext
import org.apache.spark.graphx._

object GraphX_Pregel {
  def main(args: Array[String]): Unit = {
    // 创建 SparkContext
    val sc = new SparkContext("local", "GraphPregelExample")

    // 创建一个简单的图
    val vertices = sc.parallelize(Array((1L, 0), (2L, 0), (3L, 0), (4L, 0)))
    val edges = sc.parallelize(Array(
      Edge(1L, 2L, 1), // 边1: 1 -> 2，权重为1
      Edge(2L, 3L, 1), // 边2: 2 -> 3，权重为1
      Edge(3L, 4L, 1), // 边3: 3 -> 4，权重为1
      Edge(4L, 1L, 1) // 边4: 4 -> 1，权重为1
    ))

    val graph: Graph[Int, Int] = Graph(vertices, edges, 0)

    // 定义顶点程序函数，将每个顶点的属性增加 1
    def vertexProgram(vertexId: VertexId, value: Int, message: Int): Int = value + 1

    // 定义消息发送函数，发送每个顶点的属性值
    def sendMessage(edge: EdgeTriplet[Int, Int]): Iterator[(VertexId, Int)] = Iterator((edge.dstId, edge.srcAttr))

    // 定义消息合并函数，取最大值
    def mergeMessage(msg1: Int, msg2: Int): Int = math.max(msg1, msg2)

    // 执行 Pregel 迭代计算
    val resultGraph = graph.pregel(initialMsg = 0, maxIterations = 10, activeDirection = EdgeDirection.Out)(
      vprog = vertexProgram,
      sendMsg = sendMessage,
      mergeMsg = mergeMessage
    )

    // 输出结果
    resultGraph.vertices.collect().foreach(println)
  }
}
